{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Evolutive equilibrium calculations\n",
    "\n",
    "This notebook will demonstrate the core use-case for FreeGSNKE: simulating the (time-dependent) evolution of Grad-Shafranov (GS) equilibria. In particular we will simulate a **vertical displacement event (VDE)**.\n",
    "\n",
    "To do this, we need to:\n",
    "- build the tokamak machine.\n",
    "- instatiate a GS equilibrium (to be used as an initial condition for the evolutive solver).\n",
    "- calculate a vertical instability growth rate for this equilibrium and carry out passive structure mode removal via a normal mode decomposition (i.e. removing modes that have little effect on the evolution). More details on this in a later notebook. \n",
    "- define time-dependent coil voltages, plasma current density profile parameters, and plasma resistivity.\n",
    "- evolve the active coil currents, the total plasma current, and the equilbirium using these profile parameters and voltages by solving the circuit equations alongside the GS equation.\n",
    "\n",
    "Refer to the paper by [Amorisco et al. (2024)](https://pubs.aip.org/aip/pop/article/31/4/042517/3286904/FreeGSNKE-A-Python-based-dynamic-free-boundary) for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display, clear_output\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diverted plasma example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate an equilibrium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import equilibrium_update\n",
    "\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,\n",
    "    Rmin=0.1, Rmax=2.0,   # Radial range\n",
    "    Zmin=-2.2, Zmax=2.2,  # Vertical range\n",
    "    nx=65,                # Number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,               # Number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    # psi=plasma_psi\n",
    ")  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instatiate a profile object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,        # equilibrium object\n",
    "    paxis=8.1e3,  # profile object\n",
    "    Ip=6.2e5,     # plasma current\n",
    "    fvac=0.5,     # fvac = rB_{tor}\n",
    "    alpha_m=1.8,  # profile function parameter\n",
    "    alpha_n=1.2   # profile function parameter\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set coil currents\n",
    "Here we set coil currents that create a diverted plasma (as seen in prior notebooks). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('data/simple_diverted_currents_PaxisIp.pk', 'rb') as f:\n",
    "    current_values = pickle.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq.tokamak.set_coil_current(coil_label=key, current_value=current_values[key])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instatiate the solver\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call forward solver to find equilibrium "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-8,\n",
    "                     verbose=0\n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the initial equilibrium "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time evolution\n",
    "\n",
    "We are now ready to solve the forward time-evolutive problem.\n",
    "\n",
    "#### Initialise nonlinear (evolutive) solver\n",
    "To start, we need to initialise the evolutive solver object, `freegsnke.nonlinear_solve.nl_solver`.\n",
    "\n",
    "The time-evolutive solver requires:\n",
    "- `eq`: an equilibrium to inform the solver of the machine and domain properties.\n",
    "- `profiles`: defined above.\n",
    "- `full_timestep`: the time step by which time is advanced at every call of the stepper (this can be modified later depending on the growth rate of the equilibrium). An appropriate time step can also be set based on the growth rate calculation. Use `automatic_timestep` to set the time step in units of the (inverse of the) growth rate.\n",
    "- `plasma_resistivity`: resistivity of the plasma (which here is assumed to be constant during the time evolution but can be made time-dependent if known).\n",
    "- `min_dIy_dI`: threshold value below which passive structure normal modes are dropped. Modes with norm(d(Iy)/dI)<`min_dIy_dI` are dropped, which filters out modes that do not actually couple with the plasma.\n",
    "- `max_mode_frequency`: threshold value for characteristic frequencies above which passive structure normal modes are dropped (i.e. the fast modes).\n",
    "\n",
    "Other customisable inputs are available, do see the documentation or the later notebook on \"Growth Rates\" for more details. For example, one may explicitly set your own resistance and inductance matries for the tokamak, rather than the geometrical value calculated internally in FreeGSNKE.\n",
    "\n",
    "The solver can be used on different equilibria and/or profiles, but these need to have the same machine, domain, and limiter as the one used at the time of the solver instantiation. For different machines, a new time-evolutive solver should be created.\n",
    "\n",
    "The input equilibrium and profile functions are also used as the expansion point around which the dynamics are linearised (when using the linear solver or calculating linear growth rates). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import nonlinear_solve\n",
    "\n",
    "stepping = nonlinear_solve.nl_solver(\n",
    "    eq=eq, \n",
    "    profiles=profiles, \n",
    "    GSStaticSolver=GSStaticSolver,\n",
    "    full_timestep=5e-4, \n",
    "    plasma_resistivity=1e-6,\n",
    "    max_mode_frequency=10**2.5,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set active coil voltages\n",
    "\n",
    "In this example, we will evolve a plasma in absence of any control policy or current drive.\n",
    "\n",
    "Just as an example, the following calculates active voltages to be applied to the poloidal field coils (and Solenoid) using $V = RI$, with current values as defined by the initial equilibrium (i.e. we have constant voltages).\n",
    "\n",
    "In most FreeGSNKE use cases, these active voltages will be determined by a control policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, the solver is prepared by setting the initial conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stepping.initialize_from_ICs(eq, profiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set time steps and storage lists\n",
    "Now we set the total number of time steps we want to simulate and initialise some variables for logging the evolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of time steps to simulate\n",
    "max_count = 50\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times = [t]\n",
    "history_currents = [stepping.currents_vec]\n",
    "history_equilibria = [stepping.eq1.create_auxiliary_equilibrium()]\n",
    "history_o_points = [stepping.eq1.opt[0]]\n",
    "history_elongation = [stepping.eq1.geometricElongation()]\n",
    "history_triangularity = [stepping.eq1.triangularity()]\n",
    "history_squareness = [stepping.eq1.squareness()[1]]\n",
    "history_area = [stepping.eq1.separatrix_area()]\n",
    "history_length = [stepping.eq1.separatrix_length()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set time-dependent active coil voltages, profile parameters, and plasma resistivity\n",
    "\n",
    "To simulate the VDE in the example, we will evolve a plasma in absence of any control policy or current drive. The following calculates active voltages to be applied to the poloidal field coils (and Solenoid) using $V = RI$, with current values as defined by the initial equilibrium (i.e. we have constant voltages over time). For more realistic plasma simulation, these active voltages would be truly time-dependent and determined by a control policy.\n",
    "\n",
    "In this example, the profile parameters and the plasma resistivity will also remain constant over time. In the next example, we will show how to vary these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the solver (linear)\n",
    "Finally, we call the time-evolutive solver with `stepping.nlstepper()` sequentially until we reach the preset end time.\n",
    "\n",
    "The following demonstrates a solely linear evolution of the plasma by setting `linear_only=True`. This makes use of the Jacobians calculated when the nonlinear solver object was initialised. Note that the linear solution will only hold for equilibrium \"close to\" that which was used to initialise the nonlinear solver object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise the solver with the initial equilibrium/profiles\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# loop over time steps\n",
    "while counter<max_count:\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "\n",
    "    # carry out the time step (profile parameters and resistivity are constant at each time step)\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages,   # same voltages used at each time step\n",
    "        linear_only=True,\n",
    "        verbose=False,\n",
    "        max_solving_iterations=50,\n",
    "    )\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times.append(t)\n",
    "    counter += 1\n",
    "\n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents.append(stepping.currents_vec)\n",
    "    history_equilibria.append(stepping.eq1.create_auxiliary_equilibrium())\n",
    "    history_o_points.append(stepping.eq1.opt[0])\n",
    "    history_elongation.append(stepping.eq1.geometricElongation())\n",
    "    history_triangularity.append(stepping.eq1.triangularity())\n",
    "    history_squareness.append(stepping.eq1.squareness()[1])\n",
    "    history_area.append(stepping.eq1.separatrix_area())\n",
    "    history_length.append(stepping.eq1.separatrix_length())\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents = np.array(history_currents)\n",
    "history_times = np.array(history_times)\n",
    "history_o_points = np.array(history_o_points)\n",
    "history_elongation = np.array(history_elongation)\n",
    "history_triangularity = np.array(history_triangularity)\n",
    "history_squareness = np.array(history_squareness)\n",
    "history_area = np.array(history_area)\n",
    "history_length = np.array(history_length)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the solver (nonlinear)\n",
    "\n",
    "Now we evolve the same plasma according to the full nonlinear dynamics for the same time interval. This is done by using `linear_only=False` in the call to the stepper.\n",
    "\n",
    "We need to re-initialise from the initial conditions and reset the counter, but otherwise the method is identical to the one above.\n",
    "\n",
    "Note that the full nonlinear evolutive solve is a lot more computationally expensive than solving the linear evolutive problem. As such, the following cell may take several minutes to execute, depending on your hardware."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times_nl = [t]\n",
    "history_currents_nl = [stepping.currents_vec]\n",
    "history_equilibria_nl = [stepping.eq1.create_auxiliary_equilibrium()]\n",
    "history_o_points_nl = [stepping.eq1.opt[0]]\n",
    "history_elongation_nl = [stepping.eq1.geometricElongation()]\n",
    "history_triangularity_nl = [stepping.eq1.triangularity()]\n",
    "history_squareness_nl = [stepping.eq1.squareness()[1]]\n",
    "history_area_nl = [stepping.eq1.separatrix_area()]\n",
    "history_length_nl = [stepping.eq1.separatrix_length()]\n",
    "\n",
    "# loop over the time steps\n",
    "while counter<max_count:\n",
    "\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "    \n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages, \n",
    "        linear_only=False,\n",
    "        verbose=False,\n",
    "        max_solving_iterations=50,\n",
    "    )\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times_nl.append(t)\n",
    "    counter += 1\n",
    "    \n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents_nl.append(stepping.currents_vec)\n",
    "    history_equilibria_nl.append(stepping.eq1.create_auxiliary_equilibrium())\n",
    "    history_o_points_nl.append(stepping.eq1.opt[0])\n",
    "    history_elongation_nl.append(stepping.eq1.geometricElongation())\n",
    "    history_triangularity_nl.append(stepping.eq1.triangularity())\n",
    "    history_squareness_nl.append(stepping.eq1.squareness()[1])\n",
    "    history_area_nl.append(stepping.eq1.separatrix_area())\n",
    "    history_length_nl.append(stepping.eq1.separatrix_length())\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_nl = np.array(history_currents_nl)\n",
    "history_times_nl = np.array(history_times_nl)\n",
    "history_o_points_nl = np.array(history_o_points_nl)\n",
    "history_elongation_nl = np.array(history_elongation_nl)\n",
    "history_triangularity_nl = np.array(history_triangularity_nl)\n",
    "history_squareness_nl = np.array(history_squareness_nl)\n",
    "history_area_nl = np.array(history_area_nl)\n",
    "history_length_nl = np.array(history_length_nl)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot some time-evolving quantities\n",
    "The following plots the evolution of a number of tracked values and compares the linear/nonlinear evoltuions against one another. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot evolution of tracked values and compare between linear and non-linear evolution\n",
    "fig, axs = plt.subplots(3, 3, figsize=(15, 10), dpi=80, constrained_layout=True)\n",
    "axs_flat = axs.flat\n",
    "\n",
    "axs_flat[0].plot(history_times, history_o_points[:, 0],'k+', label='linear')\n",
    "axs_flat[0].plot(history_times_nl, history_o_points_nl[:, 0],'rx', label='nonlinear')\n",
    "axs_flat[0].set_xlabel('Time')\n",
    "axs_flat[0].set_ylabel('Magnetic axis, $R$ [$m$]')\n",
    "axs_flat[0].legend()\n",
    "axs_flat[0].grid()\n",
    "\n",
    "axs_flat[1].plot(history_times, abs(history_o_points[:, 1]),'k+')\n",
    "axs_flat[1].plot(history_times_nl, abs(history_o_points_nl[:, 1]),'rx')\n",
    "axs_flat[1].set_yscale('log')\n",
    "axs_flat[1].set_xlabel('Time')\n",
    "axs_flat[1].set_ylabel('abs(Magnetic axis, $Z$) [$m$]')\n",
    "axs_flat[1].grid()\n",
    "\n",
    "axs_flat[2].plot(history_times, history_o_points[:, 2],'k+')\n",
    "axs_flat[2].plot(history_times_nl, history_o_points_nl[:, 2],'rx')\n",
    "axs_flat[2].set_xlabel('Time')\n",
    "axs_flat[2].set_ylabel('Magnetic axis, $\\psi$')\n",
    "axs_flat[2].grid()\n",
    "\n",
    "axs_flat[3].plot(history_times, history_currents[:,-1]*stepping.plasma_norm_factor,'k+')\n",
    "axs_flat[3].plot(history_times_nl, history_currents_nl[:,-1]*stepping.plasma_norm_factor,'rx')\n",
    "axs_flat[3].set_xlabel('Time')\n",
    "axs_flat[3].set_ylabel('Plasma current')\n",
    "axs_flat[3].grid()\n",
    "\n",
    "axs_flat[4].plot(history_times, history_elongation,'k+')\n",
    "axs_flat[4].plot(history_times_nl, history_elongation_nl,'rx')\n",
    "axs_flat[4].set_xlabel('Time')\n",
    "axs_flat[4].set_ylabel('Geometric elongation')\n",
    "axs_flat[4].grid()\n",
    "\n",
    "axs_flat[5].plot(history_times, history_triangularity,'k+')\n",
    "axs_flat[5].plot(history_times_nl, history_triangularity_nl,'rx')\n",
    "axs_flat[5].set_xlabel('Time')\n",
    "axs_flat[5].set_ylabel('Triangularity')\n",
    "axs_flat[5].grid()\n",
    "\n",
    "axs_flat[6].plot(history_times, history_squareness,'k+')\n",
    "axs_flat[6].plot(history_times_nl, history_squareness_nl,'rx')\n",
    "axs_flat[6].set_xlabel('Time')\n",
    "axs_flat[6].set_ylabel('Squarenss')\n",
    "axs_flat[6].grid()\n",
    "\n",
    "axs_flat[7].plot(history_times, history_area,'k+')\n",
    "axs_flat[7].plot(history_times_nl, history_area_nl,'rx')\n",
    "axs_flat[7].set_xlabel('Time')\n",
    "axs_flat[7].set_ylabel('LCFS area')\n",
    "axs_flat[7].grid()\n",
    "\n",
    "axs_flat[8].plot(history_times, history_length,'k+')\n",
    "axs_flat[8].plot(history_times_nl, history_length_nl,'rx')\n",
    "axs_flat[8].set_xlabel('Time')\n",
    "axs_flat[8].set_ylabel('LCFS circumference')\n",
    "axs_flat[8].grid()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the equilibria at the final time step\n",
    "fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)\n",
    "\n",
    "ax1.grid(True, which='both')\n",
    "history_equilibria[-1].plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "ax1.set_title(\"Linear solve\")\n",
    "\n",
    "ax2.grid(True, which='both')\n",
    "history_equilibria_nl[-1].plot(axis=ax2, show=False)\n",
    "eq.tokamak.plot(axis=ax2, show=False)\n",
    "ax2.set_xlim(0.1, 2.15)\n",
    "ax2.set_ylim(-2.25, 2.25)\n",
    "ax2.set_title(\"Nonlinear solve\")\n",
    "\n",
    "plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Limited plasma example\n",
    "\n",
    "In this example, we examine a limiter configuration plasma and show how to use time-varying plasma profile parameters and plasma resisitivity. Note there is also the option to specify time-dependent active coil resistances. For example, if you want to switch a coil \"off\" during a simulation you can set its resistance to an arbitrarily high number. See the optional parameter `custom_active_coil_resistances`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "from freegsnke import equilibrium_update, GSstaticsolver\n",
    "\n",
    "# re-instantiate the equilibrium object\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,\n",
    "    Rmin=0.1, Rmax=2.0,\n",
    "    Zmin=-2.2, Zmax=2.2,\n",
    "    nx=65, \n",
    "    ny=129\n",
    ")  \n",
    "\n",
    "# load coil currents\n",
    "with open('data/simple_limited_currents_PaxisIp.pk', 'rb') as f:\n",
    "    current_values = pickle.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq.tokamak.set_coil_current(key, current_values[key])\n",
    "\n",
    "\n",
    "# define profiles\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,\n",
    "    paxis=6e3,\n",
    "    Ip=4e5,\n",
    "    fvac=0.5,\n",
    "    alpha_m=1.8,\n",
    "    alpha_n=1.2\n",
    ")\n",
    "\n",
    "# solve\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    \n",
    "GSStaticSolver.solve(\n",
    "    eq, \n",
    "    profiles, \n",
    "    target_relative_tolerance=1e-8, \n",
    "    verbose=False\n",
    "    )\n",
    "\n",
    "# visualise:\n",
    "# the red dashed line is the flux surface through the first X-point, however, the actual last closed flux surface \n",
    "# in limiter equilibria is displayed as a full black line. \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import nonlinear_solve\n",
    "\n",
    "# initialise the nonlinear solver object\n",
    "stepping = nonlinear_solve.nl_solver(\n",
    "    eq=eq, \n",
    "    profiles=profiles, \n",
    "    full_timestep=.5e-3, \n",
    "    plasma_resistivity=1e-6, \n",
    "    GSStaticSolver=GSStaticSolver   \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of time steps to simulate\n",
    "max_count = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# determine the (constant) active voltages to be applied at each time step\n",
    "voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we define some time-dependent plasma current density profile parameters here\n",
    "alpha_m = np.tile(profiles.alpha_m, max_count+1)\n",
    "alpha_m -= (0.1 * np.sin(0.05 * np.pi * np.arange(max_count+1))) # we add some perturbation\n",
    "\n",
    "alpha_n = np.tile(profiles.alpha_n, max_count+1)\n",
    "alpha_n += (0.1 * np.sin(0.1 * np.pi * np.arange(max_count+1))) # we add some perturbation\n",
    "\n",
    "paxis = np.tile(profiles.paxis, max_count+1)\n",
    "paxis += (0.1 * np.sin(0.01 * np.pi * np.arange(max_count+1))) # we add some perturbation\n",
    "\n",
    "\n",
    "# plot of profiles\n",
    "fig, axs = plt.subplots(3, 1, figsize=(8, 6), dpi=80, constrained_layout=True)\n",
    "axs_flat = axs.flat\n",
    "\n",
    "axs_flat[0].plot(alpha_m, label='alpha_m')\n",
    "axs_flat[0].set_xticklabels([])\n",
    "axs_flat[0].set_ylabel('Parameter value')\n",
    "axs_flat[0].legend()\n",
    "axs_flat[0].grid()\n",
    "\n",
    "axs_flat[1].plot(alpha_n, label='alpha_n')\n",
    "axs_flat[1].set_xticklabels([])\n",
    "axs_flat[1].set_ylabel('Parameter value')\n",
    "axs_flat[1].legend()\n",
    "axs_flat[1].grid()\n",
    "\n",
    "axs_flat[2].plot(paxis, label='paxis')\n",
    "axs_flat[2].set_xlabel('Time Step')\n",
    "axs_flat[2].set_ylabel('Parameter value')\n",
    "axs_flat[2].legend()\n",
    "axs_flat[2].grid()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we also define a time-dependent plasma resisitivity. \n",
    "resisitivity = np.tile(stepping.plasma_resistivity, max_count+1)\n",
    "resisitivity += (1e-7 * np.sin(0.05 * np.pi * np.arange(max_count+1))) # we add some perturbation\n",
    "\n",
    "# plot of resisitivity evolution\n",
    "fig, axs = plt.subplots(1, 1, figsize=(8, 3), dpi=80)\n",
    "axs.plot(resisitivity)\n",
    "axs.set_xlabel('Time Step')\n",
    "axs.set_ylabel('Resisitivity')\n",
    "axs.grid()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "history_times = [t]\n",
    "history_currents = [stepping.currents_vec]\n",
    "history_equilibria = [stepping.eq1.create_auxiliary_equilibrium()]\n",
    "history_o_points = [stepping.eq1.opt[0]]\n",
    "history_elongation = [stepping.eq1.geometricElongation()]\n",
    "history_triangularity = [stepping.eq1.triangularity()]\n",
    "history_squareness = [stepping.eq1.squareness()[1]]\n",
    "history_area = [stepping.eq1.separatrix_area()]\n",
    "history_length = [stepping.eq1.separatrix_length()]\n",
    "\n",
    "# simulate using the linear solver\n",
    "while counter<max_count:\n",
    "\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "    \n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec = voltages, # same voltages applied at each step\n",
    "        linear_only = True,\n",
    "        profiles_parameters = { # different at each step\n",
    "            \"alpha_m\": alpha_m[counter],\n",
    "            \"alpha_n\": alpha_n[counter],\n",
    "            \"paxis\": paxis[counter],\n",
    "            },\n",
    "        plasma_resistivity = resisitivity[counter] # different at each step\n",
    "        )\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times.append(t)\n",
    "    counter += 1\n",
    "    \n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents.append(stepping.currents_vec)\n",
    "    history_equilibria.append(stepping.eq1.create_auxiliary_equilibrium())\n",
    "    history_o_points.append(stepping.eq1.opt[0])\n",
    "    history_elongation.append(stepping.eq1.geometricElongation())\n",
    "    history_triangularity.append(stepping.eq1.triangularity())\n",
    "    history_squareness.append(stepping.eq1.squareness()[1])\n",
    "    history_area.append(stepping.eq1.separatrix_area())\n",
    "    history_length.append(stepping.eq1.separatrix_length())\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents = np.array(history_currents)\n",
    "history_times = np.array(history_times)\n",
    "history_o_points = np.array(history_o_points)\n",
    "history_elongation = np.array(history_elongation)\n",
    "history_triangularity = np.array(history_triangularity)\n",
    "history_squareness = np.array(history_squareness)\n",
    "history_area = np.array(history_area)\n",
    "history_length = np.array(history_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# recalculate the active voltages using the new currents\n",
    "U_active = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] \n",
    "\n",
    "# number of time steps to simulate\n",
    "max_count = 50\n",
    "\n",
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "history_times_nl = [t]\n",
    "history_currents_nl = [stepping.currents_vec]\n",
    "history_equilibria_nl = [stepping.eq1.create_auxiliary_equilibrium()]\n",
    "history_o_points_nl = [stepping.eq1.opt[0]]\n",
    "history_elongation_nl = [stepping.eq1.geometricElongation()]\n",
    "history_triangularity_nl = [stepping.eq1.triangularity()]\n",
    "history_squareness_nl = [stepping.eq1.squareness()[1]]\n",
    "history_area_nl = [stepping.eq1.separatrix_area()]\n",
    "history_length_nl = [stepping.eq1.separatrix_length()]\n",
    "\n",
    "# simulate using the nonlinear solver\n",
    "while counter<max_count:\n",
    "\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "    \n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec = voltages, # same voltages applied at each step\n",
    "        linear_only = False,\n",
    "        profiles_parameters = { # different profile parameters at each step\n",
    "            \"alpha_m\": alpha_m[counter],\n",
    "            \"alpha_n\": alpha_n[counter],\n",
    "            \"paxis\": paxis[counter],\n",
    "            },\n",
    "        plasma_resistivity = resisitivity[counter] # different resistivity at each step\n",
    "        )\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times_nl.append(t)\n",
    "    counter += 1\n",
    "    \n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents_nl.append(stepping.currents_vec)\n",
    "    history_equilibria_nl.append(stepping.eq1.create_auxiliary_equilibrium())\n",
    "    history_o_points_nl.append(stepping.eq1.opt[0])\n",
    "    history_elongation_nl.append(stepping.eq1.geometricElongation())\n",
    "    history_triangularity_nl.append(stepping.eq1.triangularity())\n",
    "    history_squareness_nl.append(stepping.eq1.squareness()[1])\n",
    "    history_area_nl.append(stepping.eq1.separatrix_area())\n",
    "    history_length_nl.append(stepping.eq1.separatrix_length())\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_nl = np.array(history_currents_nl)\n",
    "history_times_nl = np.array(history_times_nl)\n",
    "history_o_points_nl = np.array(history_o_points_nl)\n",
    "history_elongation_nl = np.array(history_elongation_nl)\n",
    "history_triangularity_nl = np.array(history_triangularity_nl)\n",
    "history_squareness_nl = np.array(history_squareness_nl)\n",
    "history_area_nl = np.array(history_area_nl)\n",
    "history_length_nl = np.array(history_length_nl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot evolution of tracked values and compare between linear and non-linear evolution\n",
    "fig, axs = plt.subplots(3, 3, figsize=(15, 10), dpi=80, constrained_layout=True)\n",
    "axs_flat = axs.flat\n",
    "\n",
    "axs_flat[0].plot(history_times, history_o_points[:, 0],'k+', label='linear')\n",
    "axs_flat[0].plot(history_times_nl, history_o_points_nl[:, 0],'rx', label='nonlinear')\n",
    "axs_flat[0].set_xlabel('Time')\n",
    "axs_flat[0].set_ylabel('Magnetic axis, $R$ [$m$]')\n",
    "axs_flat[0].legend()\n",
    "axs_flat[0].grid()\n",
    "\n",
    "axs_flat[1].plot(history_times, abs(history_o_points[:, 1]),'k+')\n",
    "axs_flat[1].plot(history_times_nl, abs(history_o_points_nl[:, 1]),'rx')\n",
    "axs_flat[1].set_yscale('log')\n",
    "axs_flat[1].set_xlabel('Time')\n",
    "axs_flat[1].set_ylabel('abs(Magnetic axis, $Z$) [$m$]')\n",
    "axs_flat[1].grid()\n",
    "\n",
    "axs_flat[2].plot(history_times, history_o_points[:, 2],'k+')\n",
    "axs_flat[2].plot(history_times_nl, history_o_points_nl[:, 2],'rx')\n",
    "axs_flat[2].set_xlabel('Time')\n",
    "axs_flat[2].set_ylabel('Magnetic axis, $\\psi$')\n",
    "axs_flat[2].grid()\n",
    "\n",
    "axs_flat[3].plot(history_times, history_currents[:,-1]*stepping.plasma_norm_factor,'k+')\n",
    "axs_flat[3].plot(history_times_nl, history_currents_nl[:,-1]*stepping.plasma_norm_factor,'rx')\n",
    "axs_flat[3].set_xlabel('Time')\n",
    "axs_flat[3].set_ylabel('Plasma current')\n",
    "axs_flat[3].grid()\n",
    "\n",
    "axs_flat[4].plot(history_times, history_elongation,'k+')\n",
    "axs_flat[4].plot(history_times_nl, history_elongation_nl,'rx')\n",
    "axs_flat[4].set_xlabel('Time')\n",
    "axs_flat[4].set_ylabel('Geometric elongation')\n",
    "axs_flat[4].grid()\n",
    "\n",
    "axs_flat[5].plot(history_times, history_triangularity,'k+')\n",
    "axs_flat[5].plot(history_times_nl, history_triangularity_nl,'rx')\n",
    "axs_flat[5].set_xlabel('Time')\n",
    "axs_flat[5].set_ylabel('Triangularity')\n",
    "axs_flat[5].grid()\n",
    "\n",
    "axs_flat[6].plot(history_times, history_squareness,'k+')\n",
    "axs_flat[6].plot(history_times_nl, history_squareness_nl,'rx')\n",
    "axs_flat[6].set_xlabel('Time')\n",
    "axs_flat[6].set_ylabel('Squarenss')\n",
    "axs_flat[6].grid()\n",
    "\n",
    "axs_flat[7].plot(history_times, history_area,'k+')\n",
    "axs_flat[7].plot(history_times_nl, history_area_nl,'rx')\n",
    "axs_flat[7].set_xlabel('Time')\n",
    "axs_flat[7].set_ylabel('LCFS area')\n",
    "axs_flat[7].grid()\n",
    "\n",
    "axs_flat[8].plot(history_times, history_length,'k+')\n",
    "axs_flat[8].plot(history_times_nl, history_length_nl,'rx')\n",
    "axs_flat[8].set_xlabel('Time')\n",
    "axs_flat[8].set_ylabel('LCFS circumference')\n",
    "axs_flat[8].grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the equilibria at the final time step\n",
    "fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8), dpi=80)\n",
    "\n",
    "ax1.grid(True, which='both')\n",
    "history_equilibria[-1].plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "ax1.set_title(\"Linear solve\")\n",
    "\n",
    "ax2.grid(True, which='both')\n",
    "history_equilibria_nl[-1].plot(axis=ax2, show=False)\n",
    "eq.tokamak.plot(axis=ax2, show=False)\n",
    "ax2.set_xlim(0.1, 2.15)\n",
    "ax2.set_ylim(-2.25, 2.25)\n",
    "ax2.set_title(\"Nonlinear solve\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "freegsnke_github",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}